{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f63df3e4b8d4fdba",
   "metadata": {},
   "source": [
    "# Burr Demo - observing state \n",
    "# & traveling back in time!\n",
    "\n",
    "\n",
    "<img src=\"https://github.com/user-attachments/assets/2ab9b499-7ca2-4ae9-af72-ccc775f30b4e\" width=\"100\" align=\"left\" /> + \n",
    "<img src=\"https://cdn.mos.cms.futurecdn.net/VgGxJABA8DcfAMpPPwdv6a.jpg\" width=\"200\" align=\"center\"/>\n",
    "\n",
    "[https://github.com/dagworks-inc/burr](https://github.com/dagworks-inc/burr) by DAGWorks Inc. (YCW23 & StartX).\n",
    "\n",
    "Take🏠:\n",
    "\n",
    " - high level what is Burr\n",
    " - what you can do with Burr (observing state & being able to debug a particular point in time)\n",
    " - watch a walkthrough of this [notebook here](https://youtu.be/hqutVJyd3TI)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81128a911f2cbeee",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Agentic Problems"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ff9f53c-f080-4d05-8b31-70db743f973d",
   "metadata": {},
   "source": [
    "## 1. Why did this LLM call fail?\n",
    "## 2. Oh crap my code broke, why?\n",
    "## 3. Things went off the rails, but where?\n",
    "## 4. etc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d725b4dd413d01e8",
   "metadata": {},
   "source": [
    "# Monitoring FTW, right?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92cea7b3-1d27-4820-90b5-7aba9026fdfe",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Well but ... monitoring doesn't help you debug & complete your dev loop\n",
    "\n",
    "## 1. How do I debug that quickly?\n",
    "\n",
    "## 2. How do I fix the inputs/code, and restart my agent?\n",
    "\n",
    "## 3. What if my agent was 20+ steps in … do I have to restart from step 0? or can I go to a specific point in time?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f602812d90d20d69",
   "metadata": {},
   "source": [
    "# Solution: Burr\n",
    "(Complements our other framework [Hamilton](https://github.com/dagWorks-Inc/hamilton))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702e9184-7a7e-43dc-9e4a-2b4a94bfcb2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install libs if you need them\n",
    "!pip install \"burr[start]\" openai opentelemetry-instrumentation-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3422aed8-323f-45dd-ade1-b1f0906c35f3",
   "metadata": {},
   "source": [
    "## 1. Agent application is modeled as State + Actions --> Graph\n",
    "Straightforward multi-modal example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "7953fe759ae0f510",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-08-12T05:44:44.341144Z",
     "start_time": "2024-08-12T05:44:42.889221Z"
    }
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from IPython.core.display import HTML \n",
    "import openai\n",
    "\n",
    "from burr.core import ApplicationBuilder, State, default, graph, when\n",
    "from burr.core.action import action\n",
    "from burr.tracking import LocalTrackingClient\n",
    "\n",
    "MODES = {\n",
    "    \"answer_question\": \"text\",\n",
    "    \"generate_image\": \"image\",\n",
    "    \"generate_code\": \"code\",\n",
    "    \"unknown\": \"text\",\n",
    "}\n",
    "\n",
    "\n",
    "@action(reads=[], writes=[\"chat_history\", \"prompt\"])\n",
    "def process_prompt(state: State, prompt: str) -> State:\n",
    "    result = {\"chat_item\": {\"role\": \"user\", \"content\": prompt, \"type\": \"text\"}}\n",
    "    state = state.append(chat_history=result[\"chat_item\"])\n",
    "    state = state.update(prompt=prompt)\n",
    "    return state\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\"], writes=[\"mode\"])\n",
    "def choose_mode(state: State) -> State:\n",
    "    prompt = (\n",
    "        f\"You are a chatbot. You've been prompted this: {state['prompt']}. \"\n",
    "        f\"You have the capability of responding in the following modes: {', '.join(MODES)}. \"\n",
    "        \"Please respond with *only* a single word representing the mode that most accurately \"\n",
    "        \"corresponds to the prompt. Fr instance, if the prompt is 'draw a picture of a cat', \"\n",
    "        \"the mode would be 'generate_image'. If the prompt is \"\n",
    "        \"'what is the capital of France', the mode would be 'answer_question'.\"\n",
    "        \"If none of these modes apply, please respond with 'unknown'.\"\n",
    "    )\n",
    "    \n",
    "    llm_result = openai.Client().chat.completions.create(\n",
    "        model=\"gpt-4\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "    )\n",
    "    content = llm_result.choices[0].message.content\n",
    "    mode = content.lower()\n",
    "    if mode not in MODES:\n",
    "        mode = \"unknown\"\n",
    "    result = {\"mode\": mode}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\"], writes=[\"response\"])\n",
    "def prompt_for_more(state: State) -> State:\n",
    "    result = {\n",
    "        \"response\": {\n",
    "            \"content\": \"None of the response modes I support apply to your question. \"\n",
    "                       \"Please clarify?\",\n",
    "            \"type\": \"text\",\n",
    "            \"role\": \"assistant\",\n",
    "        }\n",
    "    }\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def chat_response(\n",
    "        state: State, prepend_prompt: str, model: str = \"gpt-3.5-turbo\"\n",
    ") -> State:\n",
    "    \n",
    "    chat_history = copy.deepcopy(state[\"chat_history\"])\n",
    "    chat_history[-1][\"content\"] = f\"{prepend_prompt}: {chat_history[-1]['content']}\"\n",
    "    chat_history_api_format = [\n",
    "        {\n",
    "            \"role\": chat[\"role\"],\n",
    "            \"content\": chat[\"content\"],\n",
    "        }\n",
    "        for chat in chat_history\n",
    "    ]\n",
    "    client = openai.Client()\n",
    "    result = client.chat.completions.create(\n",
    "        model=model,\n",
    "        messages=chat_history_api_format,\n",
    "    )\n",
    "    text_response = result.choices[0].message.content\n",
    "    result = {\"response\": {\"content\": text_response, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"prompt\", \"chat_history\", \"mode\"], writes=[\"response\"])\n",
    "def image_response(state: State, model: str = \"dall-e-2\") -> State:\n",
    "    \"\"\"Generates an image response to the prompt. Optional save function to save the image to a URL.\"\"\"\n",
    "    raise ValueError(\"Demo error\")\n",
    "    client = openai.Client()\n",
    "    result = client.images.generate(\n",
    "        model=model, prompt=state[\"prompt\"], size=\"1024x1024\", quality=\"standard\", n=1\n",
    "    )\n",
    "    image_url = result.data[0].url\n",
    "    result = {\"response\": {\"content\": image_url, \"type\": MODES[state[\"mode\"]], \"role\": \"assistant\"}}\n",
    "    return state.update(**result)\n",
    "\n",
    "\n",
    "@action(reads=[\"response\", \"mode\"], writes=[\"chat_history\"])\n",
    "def response(state: State) -> State:\n",
    "    # you'd do something specific here based on prior state\n",
    "    result = {\"chat_item\": state[\"response\"]}\n",
    "    return state.append(chat_history=result[\"chat_item\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "add99d139d3d72a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:23:11.529415Z",
     "start_time": "2024-07-16T18:23:11.008552Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"563pt\" height=\"307pt\"\n",
       " viewBox=\"0.00 0.00 563.47 307.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 303)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-303 559.47,-303 559.47,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M467.1,-233.4C467.1,-233.4 426.75,-233.4 426.75,-233.4 420.75,-233.4 414.75,-227.4 414.75,-221.4 414.75,-221.4 414.75,-208.8 414.75,-208.8 414.75,-202.8 420.75,-196.8 426.75,-196.8 426.75,-196.8 467.1,-196.8 467.1,-196.8 473.1,-196.8 479.1,-202.8 479.1,-208.8 479.1,-208.8 479.1,-221.4 479.1,-221.4 479.1,-227.4 473.1,-233.4 467.1,-233.4\"/>\n",
       "<text text-anchor=\"middle\" x=\"446.93\" y=\"-209.3\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M369.6,-167.8C369.6,-167.8 290.25,-167.8 290.25,-167.8 284.25,-167.8 278.25,-161.8 278.25,-155.8 278.25,-155.8 278.25,-143.2 278.25,-143.2 278.25,-137.2 284.25,-131.2 290.25,-131.2 290.25,-131.2 369.6,-131.2 369.6,-131.2 375.6,-131.2 381.6,-137.2 381.6,-143.2 381.6,-143.2 381.6,-155.8 381.6,-155.8 381.6,-161.8 375.6,-167.8 369.6,-167.8\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-143.7\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M414.33,-196.38C401.28,-189.29 386.06,-181.01 372.17,-173.46\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"374.13,-170.54 363.67,-168.84 370.79,-176.69 374.13,-170.54\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"497.48,-299 396.38,-299 396.38,-262.4 497.48,-262.4 497.48,-299\"/>\n",
       "<text text-anchor=\"middle\" x=\"446.93\" y=\"-274.9\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M446.93,-261.98C446.93,-256.75 446.93,-250.87 446.93,-245.12\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"450.43,-245.23 446.93,-235.23 443.43,-245.23 450.43,-245.23\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M107.85,-102.2C107.85,-102.2 12,-102.2 12,-102.2 6,-102.2 0,-96.2 0,-90.2 0,-90.2 0,-77.6 0,-77.6 0,-71.6 6,-65.6 12,-65.6 12,-65.6 107.85,-65.6 107.85,-65.6 113.85,-65.6 119.85,-71.6 119.85,-77.6 119.85,-77.6 119.85,-90.2 119.85,-90.2 119.85,-96.2 113.85,-102.2 107.85,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"59.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M277.8,-136.78C239.73,-128.16 186.56,-115.98 131.21,-102.73\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"132.13,-99.35 121.59,-100.42 130.5,-106.16 132.13,-99.35\"/>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M237.73,-102.2C237.73,-102.2 150.12,-102.2 150.12,-102.2 144.12,-102.2 138.12,-96.2 138.12,-90.2 138.12,-90.2 138.12,-77.6 138.12,-77.6 138.12,-71.6 144.12,-65.6 150.12,-65.6 150.12,-65.6 237.73,-65.6 237.73,-65.6 243.73,-65.6 249.73,-71.6 249.73,-77.6 249.73,-77.6 249.73,-90.2 249.73,-90.2 249.73,-96.2 243.73,-102.2 237.73,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M292.04,-130.78C276.42,-123.48 258.14,-114.93 241.62,-107.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"243.52,-104.23 232.98,-103.16 240.55,-110.57 243.52,-104.23\"/>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M379.73,-102.2C379.73,-102.2 280.12,-102.2 280.12,-102.2 274.12,-102.2 268.12,-96.2 268.12,-90.2 268.12,-90.2 268.12,-77.6 268.12,-77.6 268.12,-71.6 274.12,-65.6 280.12,-65.6 280.12,-65.6 379.73,-65.6 379.73,-65.6 385.73,-65.6 391.73,-71.6 391.73,-77.6 391.73,-77.6 391.73,-90.2 391.73,-90.2 391.73,-96.2 385.73,-102.2 379.73,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M329.93,-130.78C329.93,-125.55 329.93,-119.67 329.93,-113.92\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"333.43,-114.03 329.93,-104.03 326.43,-114.03 333.43,-114.03\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M524.23,-102.2C524.23,-102.2 421.62,-102.2 421.62,-102.2 415.62,-102.2 409.62,-96.2 409.62,-90.2 409.62,-90.2 409.62,-77.6 409.62,-77.6 409.62,-71.6 415.62,-65.6 421.62,-65.6 421.62,-65.6 524.23,-65.6 524.23,-65.6 530.23,-65.6 536.23,-71.6 536.23,-77.6 536.23,-77.6 536.23,-90.2 536.23,-90.2 536.23,-96.2 530.23,-102.2 524.23,-102.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"472.93\" y=\"-78.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M369.76,-130.78C386.34,-123.41 405.77,-114.77 423.26,-106.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"424.21,-110.4 431.92,-103.14 421.36,-104 424.21,-110.4\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M356.48,-36.6C356.48,-36.6 303.38,-36.6 303.38,-36.6 297.38,-36.6 291.38,-30.6 291.38,-24.6 291.38,-24.6 291.38,-12 291.38,-12 291.38,-6 297.38,0 303.38,0 303.38,0 356.48,0 356.48,0 362.48,0 368.48,-6 368.48,-12 368.48,-12 368.48,-24.6 368.48,-24.6 368.48,-30.6 362.48,-36.6 356.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"329.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M120.12,-67.74C123.1,-67.01 126.04,-66.29 128.93,-65.6 180.32,-53.24 239.28,-39.74 279.99,-30.52\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"280.51,-33.99 289.49,-28.37 278.96,-27.17 280.51,-33.99\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"241.85,-167.8 146,-167.8 146,-131.2 241.85,-131.2 241.85,-167.8\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.93\" y=\"-143.7\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M156.6,-130.78C141.35,-123.55 123.54,-115.09 107.38,-107.42\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"108.95,-104.3 98.42,-103.17 105.95,-110.62 108.95,-104.3\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M193.93,-130.78C193.93,-125.55 193.93,-119.67 193.93,-113.92\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"197.43,-114.03 193.93,-104.03 190.43,-114.03 197.43,-114.03\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M231.81,-130.78C247.43,-123.48 265.71,-114.93 282.23,-107.21\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"283.3,-110.57 290.87,-103.16 280.33,-104.23 283.3,-110.57\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M231.81,-65.18C247.43,-57.88 265.71,-49.33 282.23,-41.61\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"283.3,-44.97 290.87,-37.56 280.33,-38.63 283.3,-44.97\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M329.93,-65.18C329.93,-59.95 329.93,-54.07 329.93,-48.32\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"333.43,-48.43 329.93,-38.43 326.43,-48.43 333.43,-48.43\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M433.09,-65.18C416.24,-57.69 396.44,-48.88 378.71,-41\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"380.5,-37.96 369.94,-37.1 377.66,-44.36 380.5,-37.96\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M368.88,-21.97C424.73,-26.7 523.21,-38.72 544.92,-65.6 578.23,-106.85 522.83,-159.53 482.7,-189.78\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"480.85,-186.79 474.86,-195.53 485,-192.44 480.85,-186.79\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x12891c7c0>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Built the graph.\n",
    "base_graph = (\n",
    "    graph.GraphBuilder()\n",
    "    .with_actions(\n",
    "        # these are the \"nodes\" \n",
    "        prompt=process_prompt,\n",
    "        decide_mode=choose_mode,\n",
    "        generate_image=image_response,\n",
    "        generate_code=chat_response.bind(\n",
    "            prepend_prompt=\"Please respond with *only* code and no other text (at all) to the following:\",\n",
    "        ),\n",
    "        answer_question=chat_response.bind(\n",
    "            prepend_prompt=\"Please answer the following question:\",\n",
    "        ),\n",
    "        prompt_for_more=prompt_for_more,\n",
    "        response=response,\n",
    "    )\n",
    "    .with_transitions(\n",
    "        # these are the edges between nodes, based on state.\n",
    "        (\"prompt\", \"decide_mode\", default),\n",
    "        (\"decide_mode\", \"generate_image\", when(mode=\"generate_image\")),\n",
    "        (\"decide_mode\", \"generate_code\", when(mode=\"generate_code\")),\n",
    "        (\"decide_mode\", \"answer_question\", when(mode=\"answer_question\")),\n",
    "        (\"decide_mode\", \"prompt_for_more\", default),\n",
    "        (\n",
    "            [\"generate_image\", \"answer_question\", \"generate_code\", \"prompt_for_more\"],\n",
    "            \"response\",\n",
    "        ),\n",
    "        (\"response\", \"prompt\", default),\n",
    "    )\n",
    "    .build()\n",
    ")\n",
    "base_graph.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b26ad97a43444",
   "metadata": {},
   "source": [
    "## 2. Build application --> built in checkpointing & tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "4f0bf205944272ff",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:27:12.136291Z",
     "start_time": "2024-07-16T18:27:11.518066Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Attempting to instrument while already instrumented\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"576pt\" height=\"331pt\"\n",
       " viewBox=\"0.00 0.00 575.58 330.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 326.5)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-326.5 571.58,-326.5 571.58,4 -4,4\"/>\n",
       "<!-- prompt -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>prompt</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M509.1,-254.9C509.1,-254.9 468.75,-254.9 468.75,-254.9 462.75,-254.9 456.75,-248.9 456.75,-242.9 456.75,-242.9 456.75,-230.3 456.75,-230.3 456.75,-224.3 462.75,-218.3 468.75,-218.3 468.75,-218.3 509.1,-218.3 509.1,-218.3 515.1,-218.3 521.1,-224.3 521.1,-230.3 521.1,-230.3 521.1,-242.9 521.1,-242.9 521.1,-248.9 515.1,-254.9 509.1,-254.9\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-230.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt</text>\n",
       "</g>\n",
       "<!-- decide_mode -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>decide_mode</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M433.6,-187.3C433.6,-187.3 354.25,-187.3 354.25,-187.3 348.25,-187.3 342.25,-181.3 342.25,-175.3 342.25,-175.3 342.25,-162.7 342.25,-162.7 342.25,-156.7 348.25,-150.7 354.25,-150.7 354.25,-150.7 433.6,-150.7 433.6,-150.7 439.6,-150.7 445.6,-156.7 445.6,-162.7 445.6,-162.7 445.6,-175.3 445.6,-175.3 445.6,-181.3 439.6,-187.3 433.6,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"393.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">decide_mode</text>\n",
       "</g>\n",
       "<!-- prompt&#45;&gt;decide_mode -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>prompt&#45;&gt;decide_mode</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M463.47,-218.02C452.82,-210.67 440.27,-202 428.79,-194.07\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"430.92,-191.29 420.7,-188.49 426.94,-197.05 430.92,-191.29\"/>\n",
       "</g>\n",
       "<!-- input__prompt -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__prompt</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"539.48,-322.5 438.38,-322.5 438.38,-285.9 539.48,-285.9 539.48,-322.5\"/>\n",
       "<text text-anchor=\"middle\" x=\"488.93\" y=\"-298.4\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: prompt</text>\n",
       "</g>\n",
       "<!-- input__prompt&#45;&gt;prompt -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__prompt&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M488.93,-285.62C488.93,-279.74 488.93,-273.02 488.93,-266.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"492.43,-266.74 488.93,-256.74 485.43,-266.74 492.43,-266.74\"/>\n",
       "</g>\n",
       "<!-- generate_image -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>generate_image</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M117.85,-104.2C117.85,-104.2 22,-104.2 22,-104.2 16,-104.2 10,-98.2 10,-92.2 10,-92.2 10,-79.6 10,-79.6 10,-73.6 16,-67.6 22,-67.6 22,-67.6 117.85,-67.6 117.85,-67.6 123.85,-67.6 129.85,-73.6 129.85,-79.6 129.85,-79.6 129.85,-92.2 129.85,-92.2 129.85,-98.2 123.85,-104.2 117.85,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"69.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_image</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_image -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M341.99,-167.51C269.03,-165.96 140.74,-159.76 101.68,-135.7 93.45,-130.63 86.93,-122.54 82.01,-114.39\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"85.11,-112.77 77.31,-105.6 78.93,-116.07 85.11,-112.77\"/>\n",
       "<text text-anchor=\"middle\" x=\"164.55\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_image</text>\n",
       "</g>\n",
       "<!-- generate_code -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>generate_code</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M247.73,-104.2C247.73,-104.2 160.12,-104.2 160.12,-104.2 154.12,-104.2 148.12,-98.2 148.12,-92.2 148.12,-92.2 148.12,-79.6 148.12,-79.6 148.12,-73.6 154.12,-67.6 160.12,-67.6 160.12,-67.6 247.73,-67.6 247.73,-67.6 253.73,-67.6 259.73,-73.6 259.73,-79.6 259.73,-79.6 259.73,-92.2 259.73,-92.2 259.73,-98.2 253.73,-104.2 247.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"203.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">generate_code</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;generate_code -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M342.02,-163.6C313.66,-159.45 278.75,-151.38 250.93,-135.7 240.73,-129.96 231.37,-121.31 223.72,-112.88\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"226.55,-110.81 217.4,-105.47 221.22,-115.35 226.55,-110.81\"/>\n",
       "<text text-anchor=\"middle\" x=\"309.43\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=generate_code</text>\n",
       "</g>\n",
       "<!-- answer_question -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>answer_question</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M389.73,-104.2C389.73,-104.2 290.12,-104.2 290.12,-104.2 284.12,-104.2 278.12,-98.2 278.12,-92.2 278.12,-92.2 278.12,-79.6 278.12,-79.6 278.12,-73.6 284.12,-67.6 290.12,-67.6 290.12,-67.6 389.73,-67.6 389.73,-67.6 395.73,-67.6 401.73,-73.6 401.73,-79.6 401.73,-79.6 401.73,-92.2 401.73,-92.2 401.73,-98.2 395.73,-104.2 389.73,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">answer_question</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;answer_question -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" d=\"M387.8,-150.49C383.99,-140.82 378.56,-128.85 371.93,-119.2 370.51,-117.14 368.95,-115.09 367.31,-113.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"370.06,-110.91 360.78,-105.8 364.85,-115.58 370.06,-110.91\"/>\n",
       "<text text-anchor=\"middle\" x=\"444.8\" y=\"-122.4\" font-family=\"Times,serif\" font-size=\"14.00\">mode=answer_question</text>\n",
       "</g>\n",
       "<!-- prompt_for_more -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>prompt_for_more</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M534.23,-104.2C534.23,-104.2 431.62,-104.2 431.62,-104.2 425.62,-104.2 419.62,-98.2 419.62,-92.2 419.62,-92.2 419.62,-79.6 419.62,-79.6 419.62,-73.6 425.62,-67.6 431.62,-67.6 431.62,-67.6 534.23,-67.6 534.23,-67.6 540.23,-67.6 546.23,-73.6 546.23,-79.6 546.23,-79.6 546.23,-92.2 546.23,-92.2 546.23,-98.2 540.23,-104.2 534.23,-104.2\"/>\n",
       "<text text-anchor=\"middle\" x=\"482.93\" y=\"-80.1\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">prompt_for_more</text>\n",
       "</g>\n",
       "<!-- decide_mode&#45;&gt;prompt_for_more -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>decide_mode&#45;&gt;prompt_for_more</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M445.99,-160.38C474.35,-155.05 505.36,-146.87 513.92,-135.7 519.37,-128.6 517.26,-120.53 512.21,-113.02\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"515.23,-111.19 506.16,-105.73 509.85,-115.66 515.23,-111.19\"/>\n",
       "</g>\n",
       "<!-- response -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>response</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M366.48,-36.6C366.48,-36.6 313.38,-36.6 313.38,-36.6 307.38,-36.6 301.38,-30.6 301.38,-24.6 301.38,-24.6 301.38,-12 301.38,-12 301.38,-6 307.38,0 313.38,0 313.38,0 366.48,0 366.48,0 372.48,0 378.48,-6 378.48,-12 378.48,-12 378.48,-24.6 378.48,-24.6 378.48,-30.6 372.48,-36.6 366.48,-36.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"339.93\" y=\"-12.5\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">response</text>\n",
       "</g>\n",
       "<!-- generate_image&#45;&gt;response -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>generate_image&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M130.13,-69.77C133.1,-69.03 136.05,-68.31 138.93,-67.6 190.38,-54.97 249.33,-40.84 290.02,-31.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"290.6,-34.61 299.52,-28.89 288.98,-27.8 290.6,-34.61\"/>\n",
       "</g>\n",
       "<!-- input__model -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__model</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"95.85,-187.3 0,-187.3 0,-150.7 95.85,-150.7 95.85,-187.3\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.93\" y=\"-163.2\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: model</text>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_image -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_image</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M38.84,-150.37C35.2,-140.89 32.75,-129.13 36.93,-119.2 37.79,-117.15 38.85,-115.16 40.06,-113.25\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"42.69,-115.57 46.15,-105.56 37.2,-111.23 42.69,-115.57\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;generate_code -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;generate_code</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M48.56,-150.28C49.89,-139.79 53.35,-127.06 61.93,-119.2 85.32,-97.76 101.77,-109.4 136.97,-104.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.31,-107.65 146.48,-102.34 136,-100.77 137.31,-107.65\"/>\n",
       "</g>\n",
       "<!-- input__model&#45;&gt;answer_question -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>input__model&#45;&gt;answer_question</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M96.05,-162.46C128.21,-157.75 171.03,-149.48 206.93,-135.7 220.41,-130.52 221.71,-125.03 234.93,-119.2 245.07,-114.73 256.02,-110.62 266.88,-106.94\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"267.83,-110.31 276.25,-103.88 265.66,-103.66 267.83,-110.31\"/>\n",
       "</g>\n",
       "<!-- generate_code&#45;&gt;response -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>generate_code&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M240.73,-67.15C256.75,-59.42 275.73,-50.26 292.77,-42.04\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"294.24,-45.22 301.73,-37.72 291.2,-38.92 294.24,-45.22\"/>\n",
       "</g>\n",
       "<!-- answer_question&#45;&gt;response -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>answer_question&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M339.93,-67.32C339.93,-61.44 339.93,-54.72 339.93,-48.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"343.43,-48.44 339.93,-38.44 336.43,-48.44 343.43,-48.44\"/>\n",
       "</g>\n",
       "<!-- prompt_for_more&#45;&gt;response -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>prompt_for_more&#45;&gt;response</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M444.23,-67.15C427.22,-59.34 407.05,-50.09 389,-41.81\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"390.58,-38.69 380.03,-37.7 387.66,-45.05 390.58,-38.69\"/>\n",
       "</g>\n",
       "<!-- response&#45;&gt;prompt -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>response&#45;&gt;prompt</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M378.9,-22.35C434.77,-27.62 533.29,-40.57 554.92,-67.6 589.49,-110.79 545,-174.31 514.04,-209.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"511.64,-207.01 507.54,-216.78 516.84,-211.7 511.64,-207.01\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<burr.core.application.Application at 0x12891f1f0>"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from opentelemetry.instrumentation.openai import OpenAIInstrumentor\n",
    "# this will auto instrument the openAI client. No swapping of imports required!\n",
    "OpenAIInstrumentor().instrument()\n",
    "\n",
    "tracker = LocalTrackingClient(project=\"agent-demo-otel\")\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker, \n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    # .with_identifiers(app_id=, partition_key=)\n",
    "    .with_tracker(tracker, use_otel_tracing=True)  # tracking + checkpointing; one line 🪄.\n",
    "    .build()\n",
    ")\n",
    "app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c48e44fbf36f4d",
   "metadata": {},
   "source": [
    "## 3. Comes with a UI\n",
    "View runs in the UI; Let's run the app first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48745f4c04b9d28d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? write hello world in java\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖: public class HelloWorld {\n",
      "    public static void main(String[] args) {\n",
      "        System.out.println(\"Hello, World!\");\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5efe2430b34bf8e5",
   "metadata": {},
   "source": [
    "[Open the UI - http://localhost:7241/](http://localhost:7241)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bfbd6bf99f7b14",
   "metadata": {},
   "source": [
    "## But something broke / I want to debug\n",
    "\n",
    "Use:\n",
    "\n",
    "* Application ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "f5588457208604d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "app_id = \"8903a1b8-731b-4dea-8372-b915244875e3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "b4abab5aa575628d",
   "metadata": {},
   "outputs": [],
   "source": [
    "resumed_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph)\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "    )\n",
    "    .with_tracker(tracker, use_otel_tracing=True)\n",
    "    .with_identifiers(app_id=app_id)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4b21ade9-051c-453d-b3fd-f191f970c7b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'draw me a picture of a cat', 'type': 'text'}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resumed_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "2186ef09eaff58d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://oaidalleapiprodscus.blob.core.windows.net/private/org-UENfV68Ccazb3rdeYDLRWmoF/user-KU9L39oZGkNLkg6ZfSLMc4AV/img-IV9ewtD8NMArvLd5sjSFJveG.png?st=2024-08-13T21%3A19%3A42Z&se=2024-08-13T23%3A19%3A42Z&sp=r&sv=2023-11-03&sr=b&rscd=inline&rsct=image/png&skoid=d505667d-d6c1-4a0a-bac7-5c84a87759f8&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skt=2024-08-13T21%3A52%3A55Z&ske=2024-08-14T21%3A52%3A55Z&sks=b&skv=2023-11-03&sig=UJ47LW1zV8AIbxFzZyLOtTGv0MeK6Kribr1l7mVfJrg%3D\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = resumed_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bec09d1e740a30d",
   "metadata": {},
   "source": [
    "## Actually what if I want to go back to a certain point in time?\n",
    "\n",
    "* Fork: Start with state from any checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c008089be309c098",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:49:55.910255Z",
     "start_time": "2024-07-16T18:49:55.906299Z"
    }
   },
   "outputs": [],
   "source": [
    "app_id = \"8903a1b8-731b-4dea-8372-b915244875e3\"\n",
    "sequence_id = 1\n",
    "# partition_key = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "a8c5631cfebe8e48",
   "metadata": {},
   "outputs": [],
   "source": [
    "forked_app = (\n",
    "    ApplicationBuilder()\n",
    "    .with_graph(base_graph) # this could be different...\n",
    "    .initialize_from(\n",
    "        tracker,\n",
    "        resume_at_next_action=True, \n",
    "        default_state={\"chat_history\": []},\n",
    "        default_entrypoint=\"prompt\",\n",
    "        fork_from_app_id=app_id,\n",
    "        fork_from_sequence_id=sequence_id,\n",
    "        # fork_from_partition_key=partition_key\n",
    "    )\n",
    "    .with_tracker(tracker)\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "2a725130e0532f0c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-16T18:58:04.223560Z",
     "start_time": "2024-07-16T18:58:04.218274Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'role': 'user', 'content': 'draw me a picture of a cat', 'type': 'text'}]"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# show prior forked state\n",
    "forked_app.state[\"chat_history\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "e27a17968c62aa98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"https://oaidalleapiprodscus.blob.core.windows.net/private/org-UENfV68Ccazb3rdeYDLRWmoF/user-KU9L39oZGkNLkg6ZfSLMc4AV/img-fLn3DA9gtMsyjDggHnpb8JyB.png?st=2024-08-13T21%3A22%3A41Z&se=2024-08-13T23%3A22%3A41Z&sp=r&sv=2023-11-03&sr=b&rscd=inline&rsct=image/png&skoid=d505667d-d6c1-4a0a-bac7-5c84a87759f8&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skt=2024-08-13T03%3A24%3A20Z&ske=2024-08-14T03%3A24%3A20Z&sks=b&skv=2023-11-03&sig=Ok8K/VN9qoepqU4Y%2Bm%2BIy7Z9Vj4NAZV34FJZr68tQH4%3D\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Hi, how can I help? quit\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "    user_input = input(\"Hi, how can I help?\")\n",
    "    if \"quit\" == user_input.lower():\n",
    "        break\n",
    "    last_action, action_result, app_state = forked_app.run(\n",
    "        halt_after=[\"response\"], \n",
    "        inputs={\"prompt\": user_input}\n",
    "    )\n",
    "    last_message = app_state[\"chat_history\"][-1]\n",
    "    if last_message['type'] == 'image':\n",
    "        display(Image(url=last_message[\"content\"]))\n",
    "    else:\n",
    "        print(f\"🤖: {last_message['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3eb81d1f240d926",
   "metadata": {},
   "source": [
    "# Want to know more?\n",
    "\n",
    "[Link to video walking through this notebook](https://youtu.be/hqutVJyd3TI).\n",
    "\n",
    "[https://github.com/dagworks-inc/burr](https://github.com/dagworks-inc/burr)\n",
    "<img src=\"burr_qrcode.png\" width=\"125\"/>\n",
    "\n",
    "[Time Travel blog post & video:](https://blog.dagworks.io/p/travel-back-in-time-with-burr)\n",
    "\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\"><img src=\"https://img.youtube.com/vi/98vxhIcE6NI/0.jpg\"> \n",
    "</img>\n",
    "<a href=\"https://www.youtube.com/embed/98vxhIcE6NI?si=oV1BbSUCKa1UvX5P\">\n",
    "\n",
    "More blogs @ `blog.dagworks.io` e.g. [async & streaming](https://blog.dagworks.io/p/streaming-chatbot-with-burr-fastapi)\n",
    "\n",
    "More [examples](https://github.com/DAGWorks-Inc/burr/tree/main/examples/):\n",
    "\n",
    "- e.g. [test case creation](https://burr.dagworks.io/examples/guardrails/creating_tests/)\n",
    "- e.g. [multi-agent collaboration](https://github.com/DAGWorks-Inc/burr/tree/main/examples/multi-agent-collaboration)\n",
    "\n",
    "Follow on Twitter & LinkedIn:\n",
    "\n",
    "- https://x.com/burr_framework\n",
    "- https://x.com/dagworks\n",
    "- https://x.com/stefkrawczyk\n",
    "- https://www.linkedin.com/in/skrawczyk/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c54f8b9d-5782-45d6-b0e1-a91733cd6433",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found data for action:  generate_code\n",
      "\n",
      "Writing data to file YOUR_FIXTURE_FILE.json\n",
      "\n",
      "Add the following to your test file:\n",
      "\n",
      "import pytest\n",
      "from burr.core import state\n",
      "from burr.testing import pytest_generate_tests  # noqa: F401\n",
      "# TODO: import action you're testing, i.e. import generate_code.\n",
      "\n",
      "@pytest.mark.file_name(\"YOUR_FIXTURE_FILE.json\")\n",
      "def test_generate_code(input_state, expected_state):\n",
      "    \"\"\"Function for testing the action\"\"\"\n",
      "    input_state = state.State.deserialize(input_state)\n",
      "    expected_state = state.State.deserialize(expected_state)\n",
      "    _, output_state = generate_code(input_state)  # exercise the action\n",
      "    # TODO: choose appropriate way to evaluate the output\n",
      "    # e.g. exact match, fuzzy match, LLM grade, etc.\n",
      "    # this is exact match here on all values in state\n",
      "    assert output_state == expected_state\n",
      "    # e.g.\n",
      "    # assert 'some value' in output_state[\"response\"][\"content\"]\n",
      "    # assert llm_evaluator(..., ...) == \"Y\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!burr-test-case create  \\\n",
    "  --project-name \"agent-demo-otel\" \\\n",
    "  --partition-key \"null\" \\\n",
    "  --app-id \"1f983e4e-25d3-4ed1-b985-569131603f34\" \\\n",
    "  --sequence-id 14 \\\n",
    "  --target-file-name YOUR_FIXTURE_FILE.json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b80c6a2-61c1-4bfe-9f1c-59d308c72890",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
